package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"log"
	"net"
	"os"
	"time"
    "strings"
)

// IEC104TLSClient 实现IEC 104协议并使用TLS 1.3加密的客户端
type IEC104TLSClient struct {
	conn    net.Conn
	address string
}

// NewIEC104TLSClient 创建新的IEC104 TLS客户端
func NewIEC104TLSClient(address string) *IEC104TLSClient {
	return &IEC104TLSClient{
		address: address,
	}
}

// Connect 建立TLS连接 (增强版)
func (c *IEC104TLSClient) Connect() error {
	// 加载CA证书
	caCert, err := os.ReadFile("ca.crt")
	if err != nil {
		return fmt.Errorf("读取CA证书失败: %v", err)
	}
	caCertPool := x509.NewCertPool()
	if !caCertPool.AppendCertsFromPEM(caCert) {
		return fmt.Errorf("解析CA证书失败")
	}

	// 配置TLS 1.3 (增强安全配置)
	tlsConfig := &tls.Config{
		RootCAs:    caCertPool,
		MinVersion: tls.VersionTLS13,
		MaxVersion: tls.VersionTLS13,
		CipherSuites: []uint16{
			tls.TLS_AES_256_GCM_SHA384,
			tls.TLS_CHACHA20_POLY1305_SHA256,
		},
		CurvePreferences: []tls.CurveID{tls.X25519, tls.CurveP384},
		VerifyConnection: func(cs tls.ConnectionState) error {
			if len(cs.PeerCertificates) == 0 {
				return fmt.Errorf("无对端证书")
			}
			log.Printf("连接已验证, 使用证书: %s", cs.PeerCertificates[0].Subject)
			return nil
		},
	}

	// 建立连接
	conn, err := tls.Dial("tcp", c.address, tlsConfig)
	if err != nil {
		return fmt.Errorf("TLS连接失败: %v", err)
	}
	c.conn = conn

	// 验证TLS版本
	state := conn.ConnectionState()
	if state.Version != tls.VersionTLS13 {
		conn.Close()
		return fmt.Errorf("未使用TLS 1.3, 实际版本: %x", state.Version)
	}

	log.Printf("成功建立TLS 1.3连接, 使用密码套件: %s", tls.CipherSuiteName(state.CipherSuite))
	return nil
}

// SendAPDU 发送APDU数据单元
func (c *IEC104TLSClient) SendAPDU(data []byte) error {
	if c.conn == nil {
		return fmt.Errorf("连接未建立")
	}

	// IEC 104帧格式: 启动字符(0x68) + 长度 + 控制域 + 地址 + 信息体
	frame := make([]byte, 0, len(data)+6)
	frame = append(frame, 0x68)                              // 启动字符
	frame = append(frame, byte(len(data)+4))                 // 长度
	frame = append(frame, []byte{0x00, 0x00, 0x00, 0x00}...) // 控制域和地址(简化)
	frame = append(frame, data...)

	_, err := c.conn.Write(frame)
	return err
}

// Close 关闭连接
func (c *IEC104TLSClient) Close() {
	if c.conn != nil {
		c.conn.Close()
	}
}

// IEC104TLSServer 实现IEC 104协议并使用TLS 1.3加密的服务端
type IEC104TLSServer struct {
	listener net.Listener
	port     string
	certFile string
	keyFile  string
}

// NewIEC104TLSServer 创建新的IEC104 TLS服务端
func NewIEC104TLSServer(port, certFile, keyFile string) *IEC104TLSServer {
	return &IEC104TLSServer{
		port:     port,
		certFile: certFile,
		keyFile:  keyFile,
	}
}

// Start 启动TLS服务端 (增强版)
func (s *IEC104TLSServer) Start() error {
	// 加载服务端证书和密钥
	cert, err := tls.LoadX509KeyPair(s.certFile, s.keyFile)
	if err != nil {
		return fmt.Errorf("加载证书失败: %v", err)
	}

	// 加载CA证书
	caCert, err := os.ReadFile("ca.crt")
	if err != nil {
		return fmt.Errorf("读取CA证书失败: %v", err)
	}
	caCertPool := x509.NewCertPool()
	if !caCertPool.AppendCertsFromPEM(caCert) {
		return fmt.Errorf("解析CA证书失败")
	}

	// 配置TLS 1.3 (增强安全配置)
	// 在Connect方法中修改tlsConfig
	tlsConfig := &tls.Config{
		RootCAs:      caCertPool,
		Certificates: []tls.Certificate{clientCert}, // 添加客户端证书
		MinVersion:   tls.VersionTLS13,
		MaxVersion:   tls.VersionTLS13,
		CipherSuites: []uint16{
			tls.TLS_AES_256_GCM_SHA384,
			tls.TLS_CHACHA20_POLY1305_SHA256,
		},
		CurvePreferences: []tls.CurveID{tls.X25519, tls.CurveP384},
		VerifyConnection: func(cs tls.ConnectionState) error {
			if len(cs.PeerCertificates) == 0 {
				return fmt.Errorf("无对端证书")
			}
			log.Printf("连接已验证, 使用证书: %s", cs.PeerCertificates[0].Subject)
			return nil
		},
	}
	// 在Start方法中修改tlsConfig
	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
		ClientCAs:    caCertPool,
		ClientAuth:   tls.RequireAndVerifyClientCert, // 要求客户端证书
		MinVersion:   tls.VersionTLS13,
		MaxVersion:   tls.VersionTLS13,
		CipherSuites: []uint16{
			tls.TLS_AES_256_GCM_SHA384,
			tls.TLS_CHACHA20_POLY1305_SHA256,
		},
		CurvePreferences: []tls.CurveID{tls.X25519, tls.CurveP384},
	}

	// 创建监听
	listener, err := tls.Listen("tcp", ":"+s.port, tlsConfig)
	if err != nil {
		return fmt.Errorf("监听失败: %v", err)
	}
	s.listener = listener

	go s.acceptConnections()
	log.Printf("IEC 104 TLS服务端已启动, 监听端口 %s", s.port)
	return nil
}

// handleConnection 处理连接 (实现完整状态机)
func (s *IEC104TLSServer) handleConnection(conn net.Conn) {
	defer func() {
		if err := recover(); err != nil {
			log.Printf("处理连接时发生panic: %v", err)
		}
		conn.Close()
	}()

	tlsConn, ok := conn.(*tls.Conn)
	if !ok {
		log.Printf("非TLS连接")
		return
	}

	// 执行TLS握手
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := tlsConn.HandshakeContext(ctx); err != nil {
		log.Printf("TLS握手失败: %v", err)
		return
	}

	// 验证连接状态
	state := tlsConn.ConnectionState()
	if state.Version != tls.VersionTLS13 {
		log.Printf("不支持的TLS版本: %x", state.Version)
		return
	}

	// 创建IEC 104会话
	session := NewIEC104Session(conn)
	session.Start()

	// 处理会话
	for {
		select {
		case data := <-session.recvChan:
			// 处理接收到的APDU
			if err := s.processAPDU(data); err != nil {
				log.Printf("处理APDU错误: %v", err)
				return
			}
		case err := <-session.errorChan:
			log.Printf("会话错误: %v", err)
			return
		case <-time.After(30 * time.Second):
			// 发送测试帧保持连接
			if err := session.SendTestFrame(); err != nil {
				log.Printf("发送测试帧失败: %v", err)
				return
			}
		}
	}
}

// Stop 停止服务端
func (s *IEC104TLSServer) Stop() {
	if s.listener != nil {
		s.listener.Close()
	}
}

// IEC104State 定义IEC 104协议状态
type IEC104State int

const (
	StateDisconnected IEC104State = iota
	StateConnected
	StateDataTransfer
	StateError
)

// IEC104Session 封装完整的IEC 104会话
type IEC104Session struct {
	state     IEC104State
	conn      net.Conn
	lastSent  uint16
	lastRecv  uint16
	sendChan  chan []byte
	recvChan  chan []byte
	errorChan chan error
	stopChan  chan struct{}
}

// NewIEC104Session 创建新的IEC 104会话
func NewIEC104Session(conn net.Conn) *IEC104Session {
	return &IEC104Session{
		state:     StateConnected,
		conn:      conn,
		sendChan:  make(chan []byte, 100),
		recvChan:  make(chan []byte, 100),
		errorChan: make(chan error, 10),
		stopChan:  make(chan struct{}),
	}
}

// Start 启动会话处理
func (s *IEC104Session) Start() {
	go s.sendLoop()
	go s.recvLoop()
	go s.stateMachine()
}

// 示例使用
func testIEC104TLSServer() {
	// 注意: 需要准备有效的证书和密钥文件
	server := NewIEC104TLSServer("2404", "server.crt", "server.key")

	if err := server.Start(); err != nil {
		log.Fatalf("启动服务端失败: %v", err)
	}
	defer server.Stop()

	fmt.Println("IEC 104 TLS 1.3服务端已启动")
	time.Sleep(10 * time.Minute) // 示例运行时间
}

func main() {
	// 初始化日志
	log.SetFlags(log.LstdFlags | log.Lmicroseconds | log.Lshortfile)

	// 启动服务端
	server := NewIEC104TLSServer("2404", "server.crt", "server.key")
	if err := server.Start(); err != nil {
		log.Fatalf("启动服务端失败: %v", err)
	}
	defer server.Stop()

	// 等待服务端启动
	time.Sleep(1 * time.Second)

	// 测试客户端
	client := NewIEC104TLSClient("localhost:2404")
	if err := client.Connect(); err != nil {
		log.Fatalf("客户端连接失败: %v", err)
	}
	defer client.Close()

	// 示例通信
	if err := client.SendAPDU([]byte{0x07, 0x00, 0x00, 0x00}); err != nil {
		log.Printf("发送APDU失败: %v", err)
	}
}

// acceptConnections 接受客户端连接
func (s *IEC104TLSServer) acceptConnections() {
    defer func() {
        if err := recover(); err != nil {
            log.Printf("acceptConnections发生panic: %v", err)
        }
    }()

    for {
        conn, err := s.listener.Accept()
        if err != nil {
            // 检查是否是正常关闭
            if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
                log.Printf("接受连接超时: %v", err)
                continue
            }
            
            // 检查监听器是否已关闭
            if opErr, ok := err.(*net.OpError); ok && opErr.Op == "accept" {
                if strings.Contains(err.Error(), "use of closed network connection") {
                    log.Println("监听器已关闭，停止接受连接")
                    return
                }
            }
            
            log.Printf("接受连接错误: %v", err)
            continue
        }

        // 记录新连接
        remoteAddr := conn.RemoteAddr().String()
        log.Printf("新客户端连接: %s", remoteAddr)

        // 启动goroutine处理连接
        go s.handleConnection(conn)
    }
}
